#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
数据清洗主程序
支持选择使用基础处理器或增强处理器
"""

import os
import sys
import argparse
import logging
import time
import pandas as pd

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('data_clean.log', encoding='utf-8')
    ]
)
logger = logging.getLogger('main')

# 解析命令行参数
parser = argparse.ArgumentParser(description='数据清洗处理工具')
parser.add_argument('--processor', choices=['enhanced', 'basic'], default='enhanced',
                    help='指定使用的处理器类型 (enhanced 或 basic)')
parser.add_argument('--input', type=str, default='',
                    help='输入的Excel或CSV文件路径')
parser.add_argument('--output', type=str, default='',
                    help='输出的Excel或CSV文件路径')
args = parser.parse_args()

def load_input_file(file_path):
    """
    加载输入文件
    
    Args:
        file_path: 输入文件路径
        
    Returns:
        DataFrame: 数据表
    """
    if not file_path:
        # 如果未指定文件，使用示例数据
        logger.info("未指定输入文件，使用示例数据")
        return get_sample_data()
    
    if not os.path.exists(file_path):
        logger.error(f"输入文件不存在: {file_path}")
        sys.exit(1)
    
    _, ext = os.path.splitext(file_path)
    try:
        if ext.lower() == '.csv':
            # 尝试不同编码读取CSV
            encodings = ['utf-8', 'gbk', 'gb2312']
            for encoding in encodings:
                try:
                    df = pd.read_csv(file_path, encoding=encoding)
                    logger.info(f"成功使用{encoding}编码读取CSV文件")
                    return df
                except UnicodeDecodeError:
                    continue
            
            # 如果所有编码都失败，使用默认方式
            logger.warning("无法确定正确的编码，使用系统默认编码")
            return pd.read_csv(file_path)
        
        elif ext.lower() in ['.xlsx', '.xls']:
            return pd.read_excel(file_path)
        
        else:
            logger.error(f"不支持的文件格式: {ext}")
            sys.exit(1)
    
    except Exception as e:
        logger.error(f"加载文件失败: {str(e)}")
        sys.exit(1)

def get_sample_data():
    """获取示例数据"""
    return pd.DataFrame([
        {"商品名称": "每日鲜本地油菜 约350g", "商品价格": 2.98},
        {"商品名称": "金龙鱼 大米 5kg装 新老包装随机发", "商品价格": 59.9},
        {"商品名称": "鲁花 压榨一级花生油 1.8L", "商品价格": 69.9},
        {"商品名称": "好彩头苏打饼干咸味1.08kg整箱装", "商品价格": 19.8},
        {"商品名称": "泸州老窖特曲52度500ml白酒礼盒装", "商品价格": 399.0},
        {"商品名称": "玉米猪肉馅水饺500g*2袋装", "商品价格": 25.8},
        {"商品名称": "iPhone13Pro 256G 午夜色", "商品价格": 7999.0},
        {"商品名称": "三星Galaxy S22 Ultra 5G手机", "商品价格": 6999.0},
    ])

def save_output_file(df, file_path):
    """
    保存数据到输出文件
    
    Args:
        df: 数据表
        file_path: 输出文件路径
    """
    if not file_path:
        # 如果未指定文件，生成默认文件名
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        file_path = f"cleaned_data_{timestamp}.xlsx"
    
    # 创建输出目录（如果需要）
    output_dir = os.path.dirname(file_path)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 根据文件扩展名保存
    _, ext = os.path.splitext(file_path)
    try:
        if ext.lower() == '.csv':
            df.to_csv(file_path, index=False, encoding='utf-8')
        else:
            df.to_excel(file_path, index=False)
        
        logger.info(f"数据已保存到: {file_path}")
    except Exception as e:
        logger.error(f"保存数据失败: {str(e)}")
        sys.exit(1)

def process_with_enhanced(df):
    """
    使用增强版处理器处理数据
    
    Args:
        df: 输入数据表
        
    Returns:
        DataFrame: 处理后的数据表
    """
    try:
        from backend.data_clean.enhanced_processor import EnhancedProcessor
        from backend.data_clean.enhanced_brand_recognition import BrandRecognizer
        
        # 初始化增强处理器
        logger.info("正在初始化增强版处理器...")
        
        # 初始化品牌识别器
        brands = [
            "每日鲜", "金龙鱼", "鲁花", "好彩头", "泸州老窖", 
            "iPhone", "三星", "Galaxy", "Samsung", "Apple"
        ]
        brand_recognizer = BrandRecognizer(brands=brands)
        
        # 处理数据 - 仅提取品牌并清洗
        logger.info(f"开始处理数据，共 {len(df)} 条记录...")
        start_time = time.time()
        
        # 检查必要的列是否存在
        required_columns = ['商品名称']
        for col in required_columns:
            if col not in df.columns:
                logger.error(f"数据表缺少必要的列: {col}")
                return None
        
        # 使用品牌识别器处理数据
        result_df = brand_recognizer.process_dataframe(
            df=df, 
            text_column='商品名称',
            brand_column='识别品牌',
            cleaned_text_column='清洗后文本'
        )
        
        # 将品牌信息拆分为单独的列
        result_df['品牌名称'] = result_df['识别品牌'].apply(
            lambda x: x[0]['brand'] if x and len(x) > 0 else '')
        result_df['品牌置信度'] = result_df['识别品牌'].apply(
            lambda x: x[0]['confidence'] if x and len(x) > 0 else 0.0)
        result_df['匹配类型'] = result_df['识别品牌'].apply(
            lambda x: x[0]['match_type'] if x and len(x) > 0 else '')
        
        # 移除原始品牌列表列
        del result_df['识别品牌']
        
        elapsed_time = time.time() - start_time
        logger.info(f"数据处理完成，耗时: {elapsed_time:.2f}秒")
        
        return result_df
    
    except ImportError as e:
        logger.error(f"加载增强处理器失败: {str(e)}")
        logger.info("尝试使用基础处理器...")
        return process_with_basic(df)

def process_with_basic(df):
    """
    使用基础处理器处理数据
    
    Args:
        df: 输入数据表
        
    Returns:
        DataFrame: 处理后的数据表
    """
    try:
        from backend.data_clean.processor import Processor
        
        # 初始化基础处理器
        logger.info("正在初始化基础处理器...")
        processor = Processor()
        
        # 处理数据
        logger.info(f"开始处理数据，共 {len(df)} 条记录...")
        start_time = time.time()
        
        # 检查必要的列是否存在
        required_columns = ['商品名称']
        for col in required_columns:
            if col not in df.columns:
                logger.error(f"数据表缺少必要的列: {col}")
                return None
        
        # 创建结果DataFrame
        result_df = df.copy()
        
        # 添加处理结果列
        result_df['清洗后文本'] = ''
        
        # 逐行处理
        for idx, row in df.iterrows():
            # 获取商品名称
            product_name = str(row['商品名称']) if not pd.isna(row['商品名称']) else ''
            if not product_name:
                continue
            
            # 处理文本
            processed_text = processor.process(product_name)
            
            # 更新结果
            result_df.at[idx, '清洗后文本'] = processed_text
        
        elapsed_time = time.time() - start_time
        logger.info(f"数据处理完成，耗时: {elapsed_time:.2f}秒")
        
        return result_df
    
    except ImportError as e:
        logger.error(f"加载基础处理器失败: {str(e)}")
        return None

def main():
    """主函数"""
    logger.info("===== 数据清洗处理开始 =====")
    
    # 加载输入数据
    df = load_input_file(args.input)
    if df is None or len(df) == 0:
        logger.error("无法获取有效数据，程序终止")
        sys.exit(1)
    
    # 根据参数选择处理器
    if args.processor == 'enhanced':
        logger.info("使用增强版处理器")
        result_df = process_with_enhanced(df)
    else:
        logger.info("使用基础处理器")
        result_df = process_with_basic(df)
    
    # 检查处理结果
    if result_df is None:
        logger.error("数据处理失败，程序终止")
        sys.exit(1)
    
    # 输出结果
    logger.info("数据处理结果预览:")
    for idx, row in result_df.head(3).iterrows():
        logger.info(f"记录 {idx}: 原文='{row['商品名称']}', 清洗后='{row['清洗后文本']}'")
    
    # 保存结果
    save_output_file(result_df, args.output)
    
    logger.info("===== 数据清洗处理完成 =====")

if __name__ == "__main__":
    main() 